package utils.prefixTree;


import java.util.Arrays;
import java.util.Iterator;

public class RyabkoTree implements PrefixStack {
    private int[] stack;
    private int[] tree;
    private int size = 0;

    public RyabkoTree() {
        stack = new int[1000];
    }

    public void push(int x) {
        if (size == stack.length) {
            stack = Arrays.copyOf(stack, size * 2);
        }
        stack[size++] = x;
        constructTree();
    }

    public int pop() {
        int val = set(--size, 0);
        update(size - 1, 0);
        return val;
    }

    public int get(int i) {
        return stack[i];
    }

    public int set(int i, int x) {
        int val = stack[i];
        stack[i] = x;
        update(i, x - val);
        return val;
    }

    public long prefixSum(int i) {
        long sum = 0;

        i++;

        while (i > 0) {
            sum += tree[i];
            i -= i & -i;
        }

        return sum;
    }

    public int size() {
        return size;
    }

    public Iterator<Integer> iterator() {
        return Arrays.stream(stack).iterator();
    }

    private void constructTree() {
        int size = size();
        tree = new int[size + 1];
        for (int i = 0; i < size; i++) {
            update(i, stack[i]);
        }
    }

    private void update(int index, int val) {
        index++;
        while (index <= size()) {
            tree[index] += val;
            index += index & -index;
        }
    }
}
